# -*- coding: utf-8 -*-
import numpy as np
from netCDF4 import Dataset


def GetAll(input_path, gauge_loc):
    '''calculate the mean of whole domain'''
    if input_path and gauge_loc:
        # read wrf output
        rainfall = Dataset(input_path)
        rainc = rainfall.variables['RAINC']
        rainnc = rainfall.variables['RAINNC']
        (Time_Number, SN_Number, WE_Number) = rainnc.shape
        # read row and column of gauges
        # row_col = pd.read_csv(gauge_loc, sep=',', names=['row', 'col'],
        # index_col=False, keep_default_na=False)
        row_col = np.genfromtxt(gauge_loc, delimiter=',', dtype=np.int)
        wrf_output = []
        for t in range(Time_Number):
            # get the rainfall data according to time
            rainc_2D = rainc[t, :, :]
            rainnc_2D = rainnc[t, :, :]
            precipitation = rainc_2D + rainnc_2D
            # get all 50 value
            # simulation = [precipitation[x, y] for x, y in row_col]
            simulation = []
            nrows, ncols = precipitation.shape
            for x, y in row_col:
                eight_grids = []
                if y - 1 >= 0:
                    n_value = precipitation[x, y - 1]
                    eight_grids.append(n_value)
                if x + 1 < ncols and y - 1 >= 0:
                    ne_value = precipitation[x + 1, y - 1]
                    eight_grids.append(ne_value)
                if x + 1 < ncols:
                    e_value = precipitation[x + 1, y]
                    eight_grids.append(e_value)
                if x + 1 < ncols and y + 1 < nrows:
                    se_value = precipitation[x + 1, y + 1]
                    eight_grids.append(se_value)
                if y + 1 < nrows:
                    s_value = precipitation[x, y + 1]
                    eight_grids.append(s_value)
                if x - 1 >= 0 and y + 1 < nrows:
                    sw_value = precipitation[x - 1, y + 1]
                    eight_grids.append(sw_value)
                if x - 1 >= 0:
                    w_value = precipitation[x - 1, y]
                    eight_grids.append(w_value)
                if x - 1 >= 0 and y - 1 >= 0:
                    nw_value = precipitation[x - 1, y - 1]
                    eight_grids.append(nw_value)
                simulation.append(np.mean(eight_grids))
            wrf_output.append(simulation)

            # extract the simulation value at gauge location
        rainfall.close()
        return np.asarray(wrf_output)


def GetMean(source):
    '''calculate the mean of whole domain'''
    if source:
        rainfall = Dataset(source)
        rainc = rainfall.variables['RAINC']
        rainnc = rainfall.variables['RAINNC']
        (Time_Number, SN_Number, WE_Number) = rainnc.shape
        mean = []
        for t in range(Time_Number):
            # get the rainfall data according to time
            rainc_2D = rainc[t, :, :]
            rainnc_2D = rainnc[t, :, :]
            precipitation = rainc_2D + rainnc_2D
            # it is not normal distribution
            # so not appropriate to use mean
            mean.append(np.mean(precipitation))
            # extract the simulation value at gauge location
        rainfall.close()
        return mean


def GetCenter(source):
    '''calculate the mean of whole domain'''
    if source:
        rainfall = Dataset(source)
        rainc = rainfall.variables['RAINC']
        rainnc = rainfall.variables['RAINNC']
        (Time_Number, SN_Number, WE_Number) = rainnc.shape
        centers = []
        for t in range(Time_Number):
            # get the rainfall data according to time
            rainc_2D = rainc[t, :, :]
            rainnc_2D = rainnc[t, :, :]
            precipitation = rainc_2D + rainnc_2D
            # get mean of the four cells' value as the value at the center
            top_left = precipitation[SN_Number // 2 - 1, WE_Number // 2 - 1]
            top_right = precipitation[SN_Number // 2 - 1, WE_Number // 2]
            bottom_left = precipitation[SN_Number // 2, WE_Number // 2 - 1]
            bottom_right = precipitation[SN_Number // 2, WE_Number // 2]
            center = np.mean([top_left, top_right, bottom_left, bottom_right])
            centers.append(center)
            # extract the simulation value at gauge location
        rainfall.close()
        return centers


def GetAccum(source):
    '''calculate the mean of whole domain'''
    if source:
        rainfall = Dataset(source)
        rainc = rainfall.variables['RAINC']
        rainnc = rainfall.variables['RAINNC']
        (Time_Number, SN_Number, WE_Number) = rainnc.shape
        accum = np.zeros((SN_Number, WE_Number))
        for t in range(Time_Number):
            # get the rainfall data according to time
            rainc_2D = rainc[t, :, :]
            rainnc_2D = rainnc[t, :, :]
            precipitation = rainc_2D + rainnc_2D
            accum += precipitation
            # extract the simulation value at gauge location
        rainfall.close()
        return accum


if __name__ == '__main__':
    print("Process WRF output netCDF")
